#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
import numpy as np
from geometry_msgs.msg import Pose2D, PoseStamped, Quaternion
from nav_msgs.msg import Path, Odometry, OccupancyGrid
import heapq
from scipy.interpolate import splprep, splev


class LocalPlanner(Node):
    def __init__(self):
        super().__init__('local_planner')

        # 订阅话题
        self.create_subscription(Odometry, '/odom', self.odom_callback, 10)
        self.create_subscription(Pose2D, '/goal', self.goal_callback, 10)
        self.create_subscription(OccupancyGrid, '/local_costmap', self.costmap_callback, 10)

        # 发布路径
        self.path_pub = self.create_publisher(Path, '/local_path', 10)

        # 初始化变量
        self.goal = None
        self.current_pose = None
        self.costmap = None

        self.get_logger().info("Local Planner node initialized.")

    def odom_callback(self, msg):
        # 获取当前机器人位姿
        self.current_pose = msg.pose.pose
        self.get_logger().debug(f"Current pose updated: {self.current_pose}")

        # 如果已经接收到目标点和代价地图，开始路径规划
        if self.goal and self.costmap:
            self.get_logger().info("Starting path creation as all data is received.")
            self.create_path()

    def goal_callback(self, msg):
        # 获取目标点位姿
        self.goal = msg
        self.get_logger().info(f"Goal received: {self.goal}")

    def costmap_callback(self, msg):
        # 获取代价地图
        self.costmap = msg
        self.get_logger().info("Costmap received.")

    def is_free(self, x, y):
        index = x + y * self.costmap.info.width
        free = 0 <= index < len(self.costmap.data) and self.costmap.data[index] < 50
        self.get_logger().debug(f"Check if cell ({x}, {y}) is free: {free}")
        return free

    def create_path(self):
        if not self.current_pose or not self.goal or not self.costmap:
            self.get_logger().warn("Cannot create path, missing current pose, goal, or costmap.")
            self.publish_empty_path()
            return

        # 转换当前位姿和目标位姿为网格坐标
        start_x = int((self.current_pose.position.x - self.costmap.info.origin.position.x) / self.costmap.info.resolution)
        start_y = int((self.current_pose.position.y - self.costmap.info.origin.position.y) / self.costmap.info.resolution)
        goal_x = int((self.goal.x - self.costmap.info.origin.position.x) / self.costmap.info.resolution)
        goal_y = int((self.goal.y - self.costmap.info.origin.position.y) / self.costmap.info.resolution)

        self.get_logger().info(f"Start position: ({start_x}, {start_y}), Goal position: ({goal_x}, {goal_y})")

        # 检查是否已经到达目标点
        if np.hypot(goal_x - start_x, goal_y - start_y) < 1:
            self.get_logger().info("Reached the goal!")
            self.publish_empty_path()
            return

        # 使用A*算法计算路径
        path = self.a_star(start_x, start_y, goal_x, goal_y)

        if path is None:
            self.get_logger().warn("No valid path found!")
            self.publish_empty_path()
            return

        self.get_logger().info(f"Path found with {len(path)} points.")

        # 对路径进行平滑处理
        smooth_path = self.smooth_path(path)

        # 如果平滑后的路径为空，发布空路径
        if not smooth_path:
            self.get_logger().warn("Smoothed path is empty.")
            self.publish_empty_path()
            return

        self.get_logger().info(f"Smoothed path has {len(smooth_path)} points.")

        # 将网格坐标转换回世界坐标并发布路径
        path_msg = Path()
        path_msg.header.frame_id = "odom"
        path_msg.header.stamp = self.get_clock().now().to_msg()

        for (x, y) in smooth_path:
            pose = PoseStamped()
            pose.header.frame_id = "odom"
            pose.pose.position.x = x * self.costmap.info.resolution + self.costmap.info.origin.position.x
            pose.pose.position.y = y * self.costmap.info.resolution + self.costmap.info.origin.position.y
            pose.pose.position.z = 0

            pose.pose.orientation = Quaternion(*quaternion_from_euler(0, 0, 0))
            path_msg.poses.append(pose)

        self.path_pub.publish(path_msg)
        self.get_logger().info("Path published.")

    def a_star(self, start_x, start_y, goal_x, goal_y):
        open_list = []
        heapq.heappush(open_list, (0, start_x, start_y))
        came_from = {}
        g_score = {(start_x, start_y): 0}
        f_score = {(start_x, start_y): self.heuristic(start_x, start_y, goal_x, goal_y)}

        self.get_logger().debug("Starting A* algorithm.")

        while open_list:
            _, current_x, current_y = heapq.heappop(open_list)
            self.get_logger().debug(f"Exploring node: ({current_x}, {current_y})")

            if current_x == goal_x and current_y == goal_y:
                self.get_logger().info("Goal reached by A* algorithm.")
                return self.reconstruct_path(came_from, current_x, current_y)

            for neighbor in self.get_neighbors(current_x, current_y):
                neighbor_x, neighbor_y = neighbor
                tentative_g_score = g_score[(current_x, current_y)] + self.dist_between(current_x, current_y, neighbor_x, neighbor_y)

                if (neighbor_x, neighbor_y) not in g_score or tentative_g_score < g_score[(neighbor_x, neighbor_y)]:
                    came_from[(neighbor_x, neighbor_y)] = (current_x, current_y)
                    g_score[(neighbor_x, neighbor_y)] = tentative_g_score
                    f_score[(neighbor_x, neighbor_y)] = tentative_g_score + self.heuristic(neighbor_x, neighbor_y, goal_x, goal_y)
                    heapq.heappush(open_list, (f_score[(neighbor_x, neighbor_y)], neighbor_x, neighbor_y))

        self.get_logger().warn("A* algorithm failed to find a path.")
        return None

    def get_neighbors(self, x, y):
        # 使用四个方向的邻居节点：上下左右
        neighbors = [
            (x-1, y), (x+1, y),
            (x, y-1), (x, y+1),
            (x-1, y-1), (x+1, y+1),
            (x-1, y+1), (x+1, y-1)
        ]
        valid_neighbors = [n for n in neighbors if 0 <= n[0] < self.costmap.info.width and 0 <= n[1] < self.costmap.info.height and self.is_free(n[0], n[1])]
        self.get_logger().debug(f"Valid neighbors for ({x}, {y}): {valid_neighbors}")
        return valid_neighbors

    def dist_between(self, x1, y1, x2, y2):
        distance = np.hypot(x2 - x1, y2 - y1)
        self.get_logger().debug(f"Distance between ({x1}, {y1}) and ({x2}, {y2}): {distance}")
        return distance

    def heuristic(self, x1, y1, x2, y2):
        # 使用欧几里得距离作为启发式函数，允许对角线运动
        heuristic_value = np.hypot(x2 - x1, y2 - y1)
        self.get_logger().debug(f"Heuristic from ({x1}, {y1}) to ({x2}, {y2}): {heuristic_value}")
        return heuristic_value

    def smooth_path(self, path):
        # 提取路径的x和y坐标
        x = [p[0] for p in path]
        y = [p[1] for p in path]
        self.get_logger().debug(f"Original path points: {len(x)}")

        # 使用B样条进行路径平滑处理
        tck, u = splprep([x, y], s=2.0)  # s参数控制平滑度
        u_fine = np.linspace(0, 1, num=100)  # 将路径细分为更多点
        x_smooth, y_smooth = splev(u_fine, tck)

        self.get_logger().debug("Path smoothed using B-spline.")
        return list(zip(x_smooth, y_smooth))

    def reconstruct_path(self, came_from, current_x, current_y):
        path = [(current_x, current_y)]
        while (current_x, current_y) in came_from:
            current_x, current_y = came_from[(current_x, current_y)]
            path.append((current_x, current_y))
        path.reverse()
        self.get_logger().debug(f"Reconstructed path: {path}")
        return path

    def publish_empty_path(self):
        # 发布空的Path消息
        empty_path_msg = Path()
        empty_path_msg.header.frame_id = "odom"
        empty_path_msg.header.stamp = self.get_clock().now().to_msg()
        self.path_pub.publish(empty_path_msg)
        self.get_logger().info("Published an empty path.")

    def run(self):
        self.get_logger().info("Local planner running...")
        rclpy.spin(self)


def quaternion_from_euler(roll, pitch, yaw):
    """
    Convert an Euler angle to a quaternion.
    
    Parameters:
    roll : float
    pitch : float
    yaw : float
    
    Returns:
    list of float: A quaternion representing the same orientation.
    """
    qx = np.sin(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) - np.cos(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
    qy = np.cos(roll/2) * np.sin(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.cos(pitch/2) * np.sin(yaw/2)
    qz = np.cos(roll/2) * np.cos(pitch/2) * np.sin(yaw/2) - np.sin(roll/2) * np.sin(pitch/2) * np.cos(yaw/2)
    qw = np.cos(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
    return [qx, qy, qz, qw]


def main(args=None):
    rclpy.init()
    planner = LocalPlanner()
    planner.run()
    planner.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
